import struct
import serial
import time
import os
from serial.tools import list_ports

FIRMWARE_SIZE_PER_PACK = 512
def invert_uint8(data):
    temp = 0
    for i in range(8):
        if data & (1<<i) :
            temp = temp | (1<<(7-i))
    return temp


def inver_uint16(data):
    temp = 0
    for i in range(16):
        if data & (1 << i):
            temp = temp | (1 << (15 - i))
    return temp


def crc16_modbus(buffer):
    crc_in = 0xFFFF
    poly = 0x8005
    temp = 0
    for item in buffer:
        temp = item
        temp = invert_uint8(temp)
        crc_in = crc_in ^ (temp << 8)
        for i in range(8):
            if crc_in & 0x8000:
                crc_in = (crc_in << 1) ^ poly
            else:
                crc_in = crc_in << 1
    crc_in = inver_uint16(crc_in)
    return crc_in


def print_buf(buffer):
    for i in buffer:
        print('%02X'%i + ' ', end='')
    print('')


def make_frame(ctrl, buffer) :
    prefix = b'\x55\xaa'
    ctrl = struct.pack('<H', ctrl)
    length = struct.pack('<H', len(buffer))
    if len(buffer) > 0 :
        body = ctrl+length+buffer
    else:
        body = ctrl+length
    crc16 = struct.pack('<H', crc16_modbus(body))
    return prefix+body+crc16


def check_frame(buffer) :
    prefix = b'\x55\xaa'
    length = len(buffer)
    if length < 8:
        return []
    for offset in range(length - 8 + 1):
        if prefix != buffer[offset:offset+2]:
            continue
        body_len = struct.unpack('<H', buffer[offset+4:offset+6])[0]
        crc16 = struct.pack('<H', crc16_modbus(buffer[offset+2:offset+2+4+body_len]))
        if crc16 != buffer[offset+2+4+body_len:offset+2+4+body_len+2]:
            continue
        else:
            return buffer[offset:offset+2+4+body_len+2]
    return []


def progress_bar(value, total):
    bar_length = 50
    cur = value / total * 1000
    progress = int(value / total * bar_length)
    print('\r', end='')
    print('[', end='')
    if progress == 0:
        print('>', end='')
        for i in range(bar_length-1):
            print(' ', end='')
    elif progress != bar_length:
        for i in range(progress):
            print('=', end='')
        print('>', end='')
        for i in range(bar_length-1-progress):
            print(' ', end='')
    else:
        for i in range(bar_length):
            print('=', end='')
    print(']', end='')
    print(' %3.1f%%'%(value / total * 100), end='')


def ser_send(ctrl, buf):
    send_buf = make_frame(ctrl, buf)
    ser.write(send_buf)
    return send_buf


def ser_read(timeout):
    timeout = timeout/1000 + time.time()
    temp_buf = bytes(0)
    while time.time() < timeout:
        if ser.in_waiting:
            temp_buf = temp_buf + ser.read(1)
        recv_buf = check_frame(temp_buf)
        if len(recv_buf) != 0:
            return recv_buf
    return []


def file_data_read(offset, size):
    file.seek(offset, 0)
    return file.read(size)


def get_pack_info():
    return file_data_read(file.seek(0, 2), 32)


def connect():
    ctrl = 0x0000
    file.seek(0,2)
    file_size = file.tell()
    file_size = struct.pack('<i', file_size)
    print_buf(file_size)
    ser_send(ctrl, file_size)
    recv_buf = ser_read(3000)
    if len(recv_buf) != 0:
        recv_buf = recv_buf[2:-2]
        reply_ctrl = struct.unpack('<H', recv_buf[0:2])[0]
        if reply_ctrl != (ctrl|0x8000):
            return 'timeout'
        data_len = struct.unpack('<H', recv_buf[2:4])[0]
        if data_len == 1:
            data = recv_buf[4:]
            if data == b'\xac' :
                return 'accept'
            else:
                return 'refuse'
    else:
        return 'timeout'


def send_firmware():
    ctrl = 0x0001
    file.seek(0,2)
    remain = file.tell()
    file_size = remain
    write_pos = 0
    count = 0
    file.seek(0, 0)
    print('Transfer progress:')
    progress_bar(file_size-remain, file_size)
    while True:
        if remain > FIRMWARE_SIZE_PER_PACK:
            read_size = FIRMWARE_SIZE_PER_PACK
        else:
            read_size = remain
        firmware_data = file_data_read(write_pos, read_size)
        data = struct.pack('<LH', write_pos, read_size) + firmware_data
        ser_send(ctrl, data)
        recv_buf = ser_read(3000)
        if len(recv_buf) != 0:
            recv_buf = recv_buf[2:-2]
            reply_ctrl = struct.unpack('<H', recv_buf[0:2])[0]
            if reply_ctrl != (ctrl|0x8000):
                return 'refuse'
            data_len = struct.unpack('<H', recv_buf[2:4])[0]
            if data_len == 1:
                data = recv_buf[4:]
                if data == b'\xac':
                    remain -= read_size
                    progress_bar(file_size - remain, file_size)
                    if remain == 0:
                        print('')
                        return 'accept'
                    else:
                        count = 0
                        write_pos += read_size
                        continue
                else:
                    return 'refuse'
        else:
            count += 1
            if count < 4:
                continue
            else:
                return "refuse"


def verify_firmware():
    ctrl = 0x0002
    ser_send(ctrl, [])
    recv_buf = ser_read(3000)
    if len(recv_buf) != 0:
        recv_buf = recv_buf[2:-2]
        reply_ctrl = struct.unpack('<H', recv_buf[0:2])[0]
        if reply_ctrl != (ctrl|0x8000):
            return 'timeout'
        data_len = struct.unpack('<H', recv_buf[2:4])[0]
        if data_len == 1:
            data = recv_buf[4:]
            if data == b'\xac':
                return 'accept'
            else:
                return 'refuse'
    else:
        return 'timeout'


def update_completed():
    ctrl = 0x0003
    ser_send(ctrl, [])
    recv_buf = ser_read(3000)
    if len(recv_buf) != 0:
        recv_buf = recv_buf[2:-2]
        reply_ctrl = struct.unpack('<H', recv_buf[0:2])[0]
        if reply_ctrl != (ctrl|0x8000):
            return 'timeout'
        data_len = struct.unpack('<H', recv_buf[2:4])[0]
        if data_len == 1:
            data = recv_buf[4:]
            if data == b'\xac':
                return 'accept'
            else:
                return 'refuse'
    else:
        return 'timeout'


def main():
    retry = 0
    connect_retry = 0
    state = 'CONNECT'
    print('start update ...')
    while True:
        print(state)
        if state == 'CONNECT':
            ret = connect() 
            if ret == 'accept':
                state = 'SEND_FIRMWARE'
            else:
                connect_retry += 1
                if connect_retry >= 3:
                    return 'Device cannot be connected normally.'
        elif state == 'SEND_FIRMWARE':
            ret = send_firmware()
            if ret == 'accept':
                state = 'VERIFY_FIRMWARE'
            elif ret == 'refuse':
                state = 'CONNECT'
        elif state == 'VERIFY_FIRMWARE':
            ret = verify_firmware()
            if ret == 'accept':
                retry = 0
                state = 'UPDATE_COMPLETED'
            elif ret == 'refuse':
                retry = 0
                state = 'CONNECT'
            elif ret == 'timeout':
                retry += 1
                if retry > 3:
                    retry = 0
                    state = 'CONNECT'
        elif state == 'UPDATE_COMPLETED':
            ret = update_completed()
            if ret == 'accept':
                break
            elif ret == 'refuse':
                retry = 0
                state = 'CONNECT'
            elif ret == 'timeout':
                retry += 1
                if retry > 3:
                    retry = 0
                    state = 'CONNECT'
    if connect_retry < 3:
        print('Update completed.')
    else:
        print('Update fail.')


port_list = list(list_ports.comports())
if len(port_list) == 0:
    print('No serial port available.')
    exit(0)
else:
    for i in range(len(port_list)):
        print(i, ": ", port_list[i])
    while True:
        select_port = input('Select the serial port, or enter \'quit\' to exit >')
        if select_port == 'quit':
            exit(0)
        select_port = int(select_port)
        if select_port >= len(port_list):
            print('Incorrect serial port selection.')
            continue
        else:
            break
print('The serial port you selected is', port_list[select_port])
file_path = '../bsp/'
file_list = os.listdir(file_path)
if len(file_list) == 0:
    print('No update firmware available.')
    exit(0)
else:
    for i in range(len(file_list)):
        print(i, ": ", file_list[i])
    while True:
        select_file = input('Select a file, or enter \'quit\' to exit >')
        if select_file == 'quit':
            exit(0)
        select_file = int(select_file)
        if select_file >= len(file_list):
            print('Incorrect file selection.')
            continue
        else:
            break
print('The selected file is', file_list[select_file])
file_name = file_path + file_list[select_file]
try:
    print('Open serial', port_list[select_port])
    ser = serial.Serial(port_list[select_port].device, baudrate=115200)
except:
    print('Serial port open error.')
    exit(-1)
try:
    print('Open file (', file_name, ').')
    file = open(file_name, "rb")
except:
    print('File open error.')
    ser.close()
    exit(-1)

# ser = serial.Serial('COM3', baudrate=115200)
# file = open('Output/BRH_STM32F030CC_VOC_APP_0201.bin', "rb")



if __name__ == '__main__':
    connect() 
    #main()
    # print('Close file (', file_name, ').')
    # file.close()
    # print('Close serial', port_list[select_port])
    # ser.close()
    # exit(0)
    